package main

import (
	log "github.com/sirupsen/logrus"
)

type ConnLimiter struct {
	count  int
	bucket chan int
}

func NewConnLimiter(count int) *ConnLimiter {
	c := new(ConnLimiter)
	c.count = count
	c.bucket = make(chan int, count)
	log.Info("初始化了令牌桶")
	return c
}

func (c *ConnLimiter) GetConn() bool {
	if len(c.bucket) >= c.count {
		return false
	}
	c.bucket <- 1
	log.Info("获取连接成功")
	return true
}

func (c *ConnLimiter) ReleaseConn() bool {
	_ = <-c.bucket
	log.Info("释放连接成功")
	return true
}
